import torch
from torch.utils.data import Dataset
import os
from PIL import Image

# class MyDataset(Dataset):
#     def __init__(self,root_dir,label_dir):
#         self.root_dir = root_dir
#         self.label_dir = label_dir
#         self.path = os.path.join(root_dir,label_dir)
#         self.img_path = os.listdir(self.path)
#
#     def __getitem__(self, idx):
#         img_name = self.img_path[idx]
#         img_item_path = os.path.join(self.root_dir,self.label_dir,img_name)
#         img = Image.open(img_item_path)
#         label = self.label_dir
#         return img,label
#
#     def __len__(self):
#         return len(self.img_path)
#
# root_dir = "D:\\PytorchLearn\\蜜蜂数据集\\hymenoptera_data\\train"
# ants_dataset = MyDataset(root_dir,"ants")
# bees_dataset = MyDataset(root_dir,"bees")
# train_dataset = ants_dataset + bees_dataset

class MyDateset(Dataset):
    def __init__(self,root_dir,state,label_dict=None):
        self.root_dir = root_dir
        self.state = state
        if label_dict is not None:
            self.label_dict = label_dict
        self.img_path = os.listdir(os.path.join(root_dir,state))

    def __getitem__(self, idx):
        img = Image.open(os.path.join(self.root_dir,self.state,self.img_path[idx]))
        if self.state == 'train':
            img_num =self.img_path[idx].split('.')[0]
            label = self.label_dict[img_num]
            return img,label
        else:
            return img

    def __len__(self):
        return len(self.img_path)
def read_csv_labels(fname):
    with open(fname,'r') as f:
        lines = f.readlines()[1:]
    tokens = [l.rstrip().split(',') for l in lines]
    return dict(((name,label) for name,label in tokens))

root_dir = "D:\\PytorchLearn\\cifar-10"
label_dict = read_csv_labels(os.path.join(root_dir,"trainLabels.csv"))

train_dataset = MyDateset(root_dir,'train',label_dict)

test_dataset = MyDateset(root_dir,'test')

train_iter = torch.utils.data.DataLoader(train_dataset,batch_size=8,shuffle=True)




